import numpy as np
import math
import base


def _psnr(img1, img2):
    # img1 and img2 have range [0, 255]
    img1 = img1.astype(np.float64)
    img2 = img2.astype(np.float64)
    mse = np.mean((img1 - img2)**2)
    if mse == 0:
        return float('inf')
    return 20 * math.log10(255.0 / math.sqrt(mse))


def psnr(img1, img2, mode="RGB"):
    if not img1.shape == img2.shape:
        raise ValueError('输入图片不相同')
    assert (mode in ["RGB", "Y"]), "运行模式只能是 RGB 或 Y"
    c = 1
    if (len(np.shape(img1)) == 3):
        h, w, c = np.shape(img1)
    if c == 3 and mode == "Y":
        img1, img2 = base.rgb2y(img1, img2)
    return _psnr(img1, img2)
